#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版脑状态分类服务 - 专注于睡眠监测

本服务专门为睡眠监测场景设计，简化了状态分类，只保留最核心的睡眠相关状态：
- 清醒状态 (awake)
- 困倦状态 (drowsy) 
- 浅睡眠状态 (light_sleep)
- 深睡眠状态 (deep_sleep)
- REM睡眠状态 (rem_sleep)
- 信号异常 (artifact)

作者: AI Assistant
创建时间: 2025-01-29
"""

import sys
import os
import logging
from typing import Dict, List, Optional, Tuple
from datetime import datetime

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.insert(0, project_root)

from models.brain_state_models import BrainWaveRatesInput, BrainStateOutput
from exception.base_exceptions import BrainStateException, ValidationException
from exception.exception_handler import global_exception_handler

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SimplifiedBrainStateService:
    """
    简化版脑状态分类服务
    
    专门为睡眠监测设计，提供6种核心状态分类：
    1. 清醒状态 - 日常清醒活动
    2. 困倦状态 - 准备入睡或半睡半醒
    3. 浅睡眠状态 - N1/N2浅睡眠阶段
    4. 深睡眠状态 - N3深度睡眠阶段
    5. REM睡眠状态 - 快速眼动睡眠
    6. 信号异常 - 数据质量问题
    """
    
    def __init__(self, band_mode: str = '5band', expected_data_points: Optional[int] = None):
        """
        初始化简化版脑状态分类服务
        
        Args:
            band_mode: 频域模式，支持 '5band' 或 '8band'
            expected_data_points: 期望的数据点数量
        """
        self.band_mode = band_mode
        self.expected_data_points = expected_data_points
        
        # 支持的频域
        self.SUPPORTED_BANDS_8 = ['delta', 'theta', 'low_alpha', 'high_alpha', 'low_beta', 'high_beta', 'low_gamma', 'mid_gamma']
        self.SUPPORTED_BANDS_5 = ['delta', 'theta', 'alpha', 'beta', 'gamma']
        
        logger.info(f"SimplifiedBrainStateService 初始化完成（{band_mode}模式），专注睡眠监测")
    
    # 简化的睡眠监测状态规则（5频域）
    SLEEP_MONITORING_RULES_5BAND = {
        'awake': {
            'description': '清醒状态',
            'chinese_name': '清醒',
            'conditions': {
                'alpha_min': 0.25,     # Alpha波存在（放松清醒）
                'alpha_max': 0.60,     # Alpha波上限
                'beta_min': 0.20,      # Beta波存在（专注清醒）
                'beta_max': 0.50,      # Beta波上限
                'theta_max': 0.25,     # Theta波限制（避免困倦混淆）
                'delta_max': 0.15,     # Delta波限制（避免睡眠混淆）
            },
            'priority': 3,
            'min_confidence': 0.6,
            'sleep_stage': 'wake'
        },
        'drowsy': {
            'description': '困倦状态',
            'chinese_name': '困倦',
            'conditions': {
                'theta_min': 0.35,     # Theta波主要特征（困倦标志）
                'theta_max': 0.65,     # Theta波上限
                'alpha_max': 0.25,     # Alpha波上限（区分清醒）
                'delta_min': 0.10,     # Delta波开始出现
                'delta_max': 0.25,     # Delta波上限（区分浅睡眠）
                'beta_max': 0.25,      # Beta波限制
            },
            'priority': 2,
            'min_confidence': 0.5,
            'sleep_stage': 'drowsy'
        },
        'light_sleep': {
            'description': '浅睡眠状态',
            'chinese_name': '浅睡眠',
            'conditions': {
                'delta_min': 0.25,     # Delta波增加（睡眠特征）
                'delta_max': 0.50,     # Delta波上限（区分深睡眠）
                'theta_min': 0.20,     # Theta波存在
                'theta_max': 0.50,     # Theta波上限
                'alpha_max': 0.20,     # Alpha波减少
                'beta_max': 0.18,      # Beta波很低
            },
            'priority': 2,
            'min_confidence': 0.5,
            'sleep_stage': 'light_sleep'
        },
        'deep_sleep': {
            'description': '深度睡眠状态',
            'chinese_name': '深睡眠',
            'conditions': {
                'delta_min': 0.50,     # Delta波主导（深睡眠标志）
                'theta_max': 0.20,     # Theta波很低
                'alpha_max': 0.15,     # Alpha波很低
                'beta_max': 0.10,      # Beta波极低
                'gamma_max': 0.08,     # Gamma波极低
            },
            'priority': 1,
            'min_confidence': 0.7,
            'sleep_stage': 'deep_sleep'
        },
        'rem_sleep': {
            'description': 'REM睡眠状态',
            'chinese_name': 'REM睡眠',
            'conditions': {
                'theta_min': 0.35,     # Theta波活跃（REM特征）
                'theta_max': 0.60,     # Theta波上限
                'beta_min': 0.20,      # Beta波相对活跃（类似清醒）
                'beta_max': 0.35,      # Beta波上限
                'delta_max': 0.12,     # Delta波很低（区分其他睡眠）
                'alpha_max': 0.18,     # Alpha波较低
                'gamma_min': 0.06,     # Gamma波存在（梦境活动）
            },
            'priority': 2,
            'min_confidence': 0.5,
            'sleep_stage': 'rem_sleep'
        },
        'artifact': {
            'description': '信号异常',
            'chinese_name': '信号异常',
            'conditions': {
                'gamma_min': 0.40,     # 极高Gamma波（干扰标志）
            },
            'priority': 1,  # 高优先级，确保异常能被识别
            'min_confidence': 0.8,
            'sleep_stage': 'artifact'
        }
    }
    
    # 简化的睡眠监测状态规则（8频域）
    SLEEP_MONITORING_RULES_8BAND = {
        'awake': {
            'description': '清醒状态',
            'chinese_name': '清醒',
            'conditions': {
                'delta_max': 0.15,
                'theta_max': 0.20,
                'low_alpha_min': 0.12, 'low_alpha_max': 0.35,
                'high_alpha_min': 0.10, 'high_alpha_max': 0.35,
                'low_beta_min': 0.08,
                'high_beta_min': 0.06,
                'low_gamma_min': 0.02,
                'mid_gamma_min': 0.01,
            },
            'priority': 3,
            'min_confidence': 0.6,
            'sleep_stage': 'wake'
        },
        'drowsy': {
            'description': '困倦状态',
            'chinese_name': '困倦',
            'conditions': {
                'delta_max': 0.25,
                'theta_min': 0.12,
                'low_alpha_max': 0.30,
                'high_alpha_max': 0.25,
                'low_beta_max': 0.15,
                'high_beta_max': 0.10,
                'low_gamma_max': 0.06,
                'mid_gamma_max': 0.04,
            },
            'priority': 2,
            'min_confidence': 0.5,
            'sleep_stage': 'drowsy'
        },
        'light_sleep': {
            'description': '浅睡眠状态',
            'chinese_name': '浅睡眠',
            'conditions': {
                'delta_min': 0.20, 'delta_max': 0.49,
                'theta_min': 0.08, 'theta_max': 0.25,
                'low_alpha_max': 0.20,
                'high_alpha_max': 0.15,
                'low_beta_max': 0.10,
                'high_beta_max': 0.08,
                'low_gamma_max': 0.04,
                'mid_gamma_max': 0.02,
            },
            'priority': 2,
            'min_confidence': 0.5,
            'sleep_stage': 'light_sleep'
        },
        'deep_sleep': {
            'description': '深度睡眠状态',
            'chinese_name': '深睡眠',
            'conditions': {
                'delta_min': 0.40,
                'theta_max': 0.25,
                'low_alpha_max': 0.15,
                'high_alpha_max': 0.10,
                'low_beta_max': 0.10,
                'high_beta_max': 0.08,
                'low_gamma_max': 0.05,
                'mid_gamma_max': 0.05,
            },
            'priority': 1,
            'min_confidence': 0.7,
            'sleep_stage': 'deep_sleep'
        },
        'rem_sleep': {
            'description': 'REM睡眠状态',
            'chinese_name': 'REM睡眠',
            'conditions': {
                'delta_max': 0.20,
                'theta_min': 0.12,
                'low_alpha_min': 0.08, 'low_alpha_max': 0.25,
                'high_alpha_min': 0.05, 'high_alpha_max': 0.20,
                'low_beta_min': 0.08,
                'high_beta_min': 0.05,
                'low_gamma_min': 0.02,
                'mid_gamma_min': 0.01,
            },
            'priority': 2,
            'min_confidence': 0.5,
            'sleep_stage': 'rem_sleep'
        },
        'artifact': {
            'description': '信号异常',
            'chinese_name': '信号异常',
            'conditions': {
                'low_gamma_min': 0.30,
                'mid_gamma_min': 0.25,
            },
            'priority': 1,
            'min_confidence': 0.8,
            'sleep_stage': 'artifact'
        }
    }
    
    def classify_brain_state(self, input_data: BrainWaveRatesInput) -> BrainStateOutput:
        """
        对脑电波数据进行简化的睡眠状态分类
        
        Args:
            input_data: 脑电波输入数据
            
        Returns:
            BrainStateOutput: 分类结果
        """
        try:
            # 验证输入数据
            validated_data = self._validate_input(input_data)
            
            # 计算平均比率
            avg_ratios = self._calculate_average_ratios(validated_data['relative_rates'])
            
            # 如果是8频域，需要合成为5频域进行分类
            if self.band_mode == '8band':
                synthesized_5bands = self._synthesize_to_5bands(avg_ratios)
                classification_ratios = synthesized_5bands
                rules = self.SLEEP_MONITORING_RULES_5BAND
            else:
                classification_ratios = avg_ratios
                rules = self.SLEEP_MONITORING_RULES_5BAND
            
            # 进行状态分类
            state, confidence = self._classify_sleep_state(classification_ratios, rules)
            
            # 记录分类数据
            self._log_classification_data(
                input_data.session_id, 
                classification_ratios, 
                state, 
                confidence
            )
            
            # 创建输出结果
            result = BrainStateOutput(
                session_id=input_data.session_id,
                state=state,
                confidence=confidence,
                details={
                    'chinese_name': rules[state]['chinese_name'],
                    'description': rules[state]['description'],
                    'sleep_stage': rules[state]['sleep_stage'],
                    'timestamp': datetime.now().isoformat(),
                    'brain_wave_ratios': classification_ratios,
                    'processing_info': {
                        'band_mode': self.band_mode,
                        'data_points': len(next(iter(validated_data['relative_rates'].values()))),
                        'service_type': 'simplified_sleep_monitoring'
                    }
                }
            )
            
            return result
            
        except Exception as e:
            logger.error(f"简化脑状态分类失败: {str(e)}")
            return self._create_error_output(input_data.session_id, str(e))
    
    def _validate_input(self, input_data: BrainWaveRatesInput) -> Dict[str, any]:
        """
        验证输入数据
        """
        if not input_data.relative_rates:
            raise ValidationException("相对比率数据不能为空")
        
        # 根据频域模式验证
        if self.band_mode == '8band':
            expected_bands = self.SUPPORTED_BANDS_8
        else:
            expected_bands = self.SUPPORTED_BANDS_5
        
        # 检查频域完整性
        missing_bands = set(expected_bands) - set(input_data.relative_rates.keys())
        if missing_bands:
            raise ValidationException(f"缺少频域数据: {missing_bands}")
        
        # 检查数据点数量
        data_lengths = [len(values) for values in input_data.relative_rates.values()]
        if len(set(data_lengths)) > 1:
            raise ValidationException("各频域数据点数量不一致")
        
        if self.expected_data_points and data_lengths[0] != self.expected_data_points:
            logger.warning(f"数据点数量({data_lengths[0]})与期望值({self.expected_data_points})不符")
        
        return {
            'relative_rates': input_data.relative_rates,
            'session_id': input_data.session_id
        }
    
    def _calculate_average_ratios(self, relative_rates: Dict[str, List[float]]) -> Dict[str, float]:
        """
        计算各频域的平均比率
        """
        avg_ratios = {}
        for band, values in relative_rates.items():
            if values:
                avg_ratios[band] = sum(values) / len(values)
            else:
                avg_ratios[band] = 0.0
        
        return avg_ratios
    
    def _classify_sleep_state(self, avg_ratios: Dict[str, float], rules: Dict) -> Tuple[str, float]:
        """
        基于简化规则进行睡眠状态分类
        """
        candidates = []
        
        # 遍历所有状态规则
        for state, rule in rules.items():
            conditions = rule['conditions']
            confidence = self._calculate_state_confidence(avg_ratios, conditions)
            
            # 检查最小置信度阈值
            min_confidence = rule.get('min_confidence', 0.5)
            if confidence >= min_confidence:
                candidates.append({
                    'state': state,
                    'confidence': confidence,
                    'priority': rule['priority']
                })
        
        # 如果没有候选状态，返回信号异常
        if not candidates:
            return 'artifact', 0.3
        
        # 按优先级和置信度排序
        candidates.sort(key=lambda x: (x['priority'], -x['confidence']))
        
        best_candidate = candidates[0]
        return best_candidate['state'], best_candidate['confidence']
    
    def _calculate_state_confidence(self, avg_ratios: Dict[str, float], conditions: Dict[str, float]) -> float:
        """
        计算状态置信度 - 改进版，严格要求所有条件都满足
        """
        if not conditions:
            return 0.0
        
        total_score = 0.0
        condition_count = 0
        failed_critical_condition = False
        
        for condition, threshold in conditions.items():
            parts = condition.split('_')
            band = parts[0]
            constraint_type = parts[1] if len(parts) > 1 else 'min'
            
            if band not in avg_ratios:
                continue
            
            ratio = avg_ratios[band]
            condition_count += 1
            
            if constraint_type == 'min':
                if ratio >= threshold:
                    # 满足最小值条件，计算超出程度
                    score = min(1.0, ratio / threshold)
                else:
                    # 不满足最小值条件，严重扣分
                    score = ratio / threshold * 0.3  # 大幅降低分数
                    failed_critical_condition = True
            elif constraint_type == 'max':
                if ratio <= threshold:
                    # 满足最大值条件，计算符合程度
                    score = 1.0 - (ratio / threshold * 0.3) if threshold > 0 else 1.0
                else:
                    # 超出最大值条件，严重扣分
                    score = max(0.0, 1.0 - (ratio - threshold) / threshold * 2.0)
                    failed_critical_condition = True
            else:
                score = 0.5  # 未知条件类型
            
            total_score += score
        
        if condition_count == 0:
            return 0.0
        
        # 计算平均置信度
        confidence = total_score / condition_count
        
        # 如果有关键条件失败，进一步降低置信度
        if failed_critical_condition:
            confidence *= 0.6  # 额外惩罚
        
        return max(0.0, min(1.0, confidence))
    
    def _synthesize_to_5bands(self, avg_ratios_8bands: Dict[str, float]) -> Dict[str, float]:
        """
        将8频域数据合成为5频域数据
        """
        synthesized = {
            'delta': avg_ratios_8bands.get('delta', 0.0),
            'theta': avg_ratios_8bands.get('theta', 0.0),
            'alpha': avg_ratios_8bands.get('low_alpha', 0.0) + avg_ratios_8bands.get('high_alpha', 0.0),
            'beta': avg_ratios_8bands.get('low_beta', 0.0) + avg_ratios_8bands.get('high_beta', 0.0),
            'gamma': avg_ratios_8bands.get('low_gamma', 0.0) + avg_ratios_8bands.get('mid_gamma', 0.0)
        }
        
        return synthesized
    
    def _create_error_output(self, session_id: str, error_message: str) -> BrainStateOutput:
        """
        创建错误输出
        """
        return BrainStateOutput(
            session_id=session_id,
            state='artifact',
            confidence=0.0,
            details={
                'chinese_name': '信号异常',
                'description': f'处理错误: {error_message}',
                'sleep_stage': 'artifact',
                'timestamp': datetime.now().isoformat(),
                'brain_wave_ratios': {},
                'processing_info': {
                    'error': error_message,
                    'service_type': 'simplified_sleep_monitoring'
                }
            }
        )
    
    def get_supported_sleep_states(self) -> Dict[str, str]:
        """
        获取支持的睡眠状态列表
        
        Returns:
            Dict[str, str]: 状态代码到中文名称的映射
        """
        rules = self.SLEEP_MONITORING_RULES_5BAND
        return {state: rule['chinese_name'] for state, rule in rules.items()}
    
    def get_sleep_state_rules(self) -> Dict[str, Dict]:
        """
        获取睡眠状态分类规则
        
        Returns:
            Dict[str, Dict]: 完整的状态规则
        """
        if self.band_mode == '8band':
            return self.SLEEP_MONITORING_RULES_8BAND
        else:
            return self.SLEEP_MONITORING_RULES_5BAND
    
    def _log_classification_data(self, session_id: str, avg_ratios: Dict[str, float], 
                               state: str, confidence: float):
        """
        记录分类数据用于调试
        """
        logger.info(f"简化睡眠状态分类结果:")
        logger.info(f"  会话ID: {session_id}")
        logger.info(f"  平均波段比例: {avg_ratios}")
        logger.info(f"  分类状态: {state}")
        logger.info(f"  置信度: {confidence:.3f}")
        logger.info(f"  服务模式: 简化睡眠监测({self.band_mode})")


if __name__ == "__main__":
    # 测试简化版服务
    print("简化版脑状态分类服务 - 专注睡眠监测")
    
    service = SimplifiedBrainStateService()
    
    # 显示支持的状态
    supported_states = service.get_supported_sleep_states()
    print(f"\n支持的睡眠状态: {supported_states}")
    
    # 显示状态规则
    state_rules = service.get_sleep_state_rules()
    print(f"\n状态分类规则数量: {len(state_rules)}")
    for state, rule in state_rules.items():
        print(f"  {state}: {rule['chinese_name']} - {rule['description']}")